import sqlite3
import os
import sys

def fix_sessions_table():
    """修复sessions表，添加缺失的category列"""
    # 数据库文件路径
    db_path = os.path.join(os.path.dirname(__file__), 'database.db')
    
    # 如果数据库文件不存在，创建它
    if not os.path.exists(db_path):
        print(f"创建新的数据库文件: {db_path}")
        open(db_path, 'w').close()
    
    try:
        # 连接到数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # 检查sessions表是否存在
        cursor.execute("""
            SELECT name FROM sqlite_master 
            WHERE type='table' AND name='sessions'
        """)
        table_exists = cursor.fetchone()
        
        if not table_exists:
            print("sessions表不存在，需要初始化数据库")
            # 创建表的SQL语句（包含category列）
            create_table_sql = """
            CREATE TABLE sessions (
                id VARCHAR(36) PRIMARY KEY,
                name VARCHAR(100) NOT NULL DEFAULT '默认会话',
                user_id VARCHAR(36),
                status VARCHAR(20) NOT NULL DEFAULT 'active',
                category VARCHAR(20) NOT NULL DEFAULT 'default',
                interaction_count INTEGER NOT NULL DEFAULT 0,
                research_count INTEGER NOT NULL DEFAULT 0,
                session_metadata TEXT,
                created_at TIMESTAMP NOT NULL,
                updated_at TIMESTAMP NOT NULL,
                closed_at TIMESTAMP,
                search_vector VARCHAR(1000)
            )
            """
            cursor.execute(create_table_sql)
            print("已创建包含category列的sessions表")
        else:
            # 检查category列是否存在
            cursor.execute("PRAGMA table_info(sessions)")
            columns = cursor.fetchall()
            column_names = [column[1] for column in columns]
            
            if 'category' not in column_names:
                print("添加category列到现有的sessions表")
                try:
                    cursor.execute("""
                        ALTER TABLE sessions 
                        ADD COLUMN category VARCHAR(20) DEFAULT 'default'
                    """)
                    print("成功添加category列")
                except sqlite3.OperationalError as e:
                    if "duplicate column name" in str(e).lower():
                        print("category列已存在")
                    else:
                        print(f"添加category列时出错: {e}")
                        # 尝试其他方式
                        print("可能需要重新创建表")
            else:
                print("category列已存在")
        
        conn.commit()
        conn.close()
        print("数据库修复完成")
        
    except Exception as e:
        print(f"修复数据库时出错: {e}")
        return False
    
    return True

if __name__ == "__main__":
    fix_sessions_table()